"""
Ardupilot 飞行器遥测接口
飞行器遥测接口提供了两个虚拟总线和多个端点，实现了获取飞行器遥测数据和发送指令到飞行器。

date: 2025/7/7
author: SiHeng Tang
file: plane_tele.py
copyright(c) DFSA Eastern Software
此程序不能提供任何担保 WITHOUT WARRANTY OF ANY KIND
"""
import logging
import time
from copy import deepcopy

from pymavlink import mavutil, mavwp

from async_sche import BaseTaskGen, TaskWrapper
from helpers import mav_messages
from helpers.geo_math_trans import GeoVector
from mengji import Vehicle
from mod_loader import deep_clear, ConfigHandler
from task_element import DataEndpoint, DataPack

logger = logging.getLogger(__name__)
handler = ConfigHandler(__name__)

DEFAULT_CONFIG = {
    "label": "plane_tele",
    "port": "tcp:192.168.1.126:5762",
    "base_req_rate_hz": 10,
    "dry_run": 0,
    "home_float_max": 5,
    "disable_messages": []
}


# TODO 全自动的消息处理句柄绑定

class PlaneTeleBus(BaseTaskGen):
    """
    使用 pymavlink 实现的 MAVLink2 固定翼飞行器遥测接口，谨慎使用多实例特性。
    """

    def __init__(self, config: dict):
        super().__init__(priority=TaskWrapper.IDLE_PRIORITY + 10)

        self._mav_connection = mavutil.mavlink_connection(config["port"])
        self._base_req_rate = config["base_req_rate_hz"]
        self._is_dry_run = config["dry_run"]
        self._home_float_max = config["home_float_max"]

        self._last_update_time = time.time()
        self._last_attitude_time = time.time()
        self._last_global_position_int_time = time.time()
        self._last_gps_raw_int_time = time.time()
        self._last_heartbeat_time = time.time()
        self._last_mission_current_time = time.time()
        self._last_mav_system_time = time.time()
        self._last_rc_channels_time = time.time()

        self.ptele_gps_ep = DataEndpoint()
        self.ptele_gps_raw_ep = DataEndpoint()
        self.ptele_attitude_ep = DataEndpoint()
        self.ptele_fm_ep = DataEndpoint()
        self.ptele_mission_current_ep = DataEndpoint()
        self.ptele_mav_system_time = DataEndpoint()
        self.ptele_rc_channels = DataEndpoint()

        self._home_position_dict = {}
        self._home_position_fixed = False
        self._mission_ack_dict = {}

        self._send_base_request()
        self._send_home_position_request()

        if self._is_dry_run:
            logger.warning(f"Plane telemetry running in dry run mode, not sending command to AP!")

    ######################################
    # 数据包处理回调，PT 只处理很少一部分遥测
    ######################################
    def _handle_msg_heartbeat(self, msg):
        # HEARTBEAT 需要有一个判断，如果没有设置 MAV_MODE_FLAG_CUSTOM_MODE_ENABLED
        # 那么这个包不可信，因为 Ardupilot 一定使用了这个位
        d = mav_messages.parse_mav_heartbeat(msg)
        if d["MAV_MODE_FLAG_CUSTOM_MODE_ENABLED"]:
            self.ptele_fm_ep.set_dp(DataPack(payload=d,
                                             interval_time=time.time() - self._last_heartbeat_time))

            self._last_heartbeat_time = time.time()

    def _handle_msg_global_position_int(self, msg):
        self.ptele_gps_ep.set_dp(DataPack(payload=mav_messages.parse_mav_global_position(msg),
                                          interval_time=time.time() - self._last_global_position_int_time))

        self._last_global_position_int_time = time.time()

    def _handle_msg_gps_raw_int(self, msg):
        self.ptele_gps_raw_ep.set_dp(DataPack(payload=mav_messages.parse_mav_gps_raw_int(msg),
                                              interval_time=time.time() - self._last_gps_raw_int_time))

        self._last_gps_raw_int_time = time.time()

    def _handle_msg_attitude(self, msg):
        self.ptele_attitude_ep.set_dp(DataPack(payload=mav_messages.parse_mav_attitude(msg),
                                               interval_time=time.time() - self._last_attitude_time))

        self._last_attitude_time = time.time()

    def _handle_msg_mission_current(self, msg):
        self.ptele_mission_current_ep.set_dp(DataPack(payload=mav_messages.parse_mav_mission_current(msg),
                                                      interval_time=time.time() - self._last_mission_current_time))

        self._last_mission_current_time = time.time()

    def _handle_msg_home_position(self, msg):
        # HOME 点监测的行为：如果 HOME 点变化，则认为 HOME 不稳定，不会返回 HOME
        # 当 HOME 连续一段时间在一定范围内，才认为 HOME 收敛
        self._home_position_dict = home_message_dict = mav_messages.parse_mav_home_position(msg)
        if self._home_position_dict:
            last_home_gv = GeoVector(lat=self._home_position_dict["latitude"],
                                     lon=self._home_position_dict["longitude"])
            this_home_gv = GeoVector(lat=home_message_dict["latitude"],
                                     lon=home_message_dict["longitude"])
            len_diff = this_home_gv.diff(last_home_gv).length()
            if len_diff < self._home_float_max:
                self._home_position_fixed = True
                logger.debug(f"home point fixed at {this_home_gv=}")
            else:
                # HOME 未收敛
                self._home_position_fixed = False
                self._send_home_position_request()
                logger.debug(f"floating home position at {this_home_gv=}")
        else:
            # 没有收敛就继续请求，有可能自动驾驶仪认为收敛或者意外解锁就不主动发了
            self._home_position_fixed = False
            self._send_home_position_request()

    def _handle_msg_mission_ack(self, msg):
        self._mission_ack_dict = mav_messages.parse_mav_mission_ack(msg)

    def _handle_msg_system_time(self, msg):
        self.ptele_mav_system_time.set_dp(DataPack(payload=mav_messages.parse_mav_system_time(msg),
                                                   interval_time=time.time() - self._last_mav_system_time))

        self._last_mav_system_time = time.time()

    def _handle_msg_rc_channels(self, msg):
        self.ptele_rc_channels.set_dp(DataPack(payload=mav_messages.parse_rc_channels(msg),
                                               interval_time=time.time() - self._last_rc_channels_time))

        self._last_rc_channels_time = time.time()

    ######################################
    # 生成器固定函数
    ######################################
    def _on_gen(self):
        # if (time.time() - self.last_update_time) < (1.0 / (3 * self.base_req_rate)):
        #     raise ExitGenerationError()
        # 这地方不用管刷新率了，直接最大限度解析，不然任务多了解析速度更不上，出现假失联的情况
        self._last_update_time = time.time()

    async def _func(self):
        """
        收集，分类，处理遥测数据
        """
        msg = self._mav_connection.recv_match(blocking=True, timeout=1.0 / (3 * self._base_req_rate))

        if not msg:
            return

        if msg.get_type() == "ATTITUDE":
            self._handle_msg_attitude(msg)

        elif msg.get_type() == "GLOBAL_POSITION_INT":
            self._handle_msg_global_position_int(msg)

        elif msg.get_type() == "GPS_RAW_INT":
            self._handle_msg_gps_raw_int(msg)

        elif msg.get_type() == "HEARTBEAT":
            self._handle_msg_heartbeat(msg)

        elif msg.get_type() == "MISSION_CURRENT":
            self._handle_msg_mission_current(msg)

        elif msg.get_type() == "HOME_POSITION":
            self._handle_msg_home_position(msg)

        elif msg.get_type() == "MISSION_ACK":
            self._handle_msg_mission_ack(msg)

        elif msg.get_type() == "SYSTEM_TIME":
            self._handle_msg_system_time(msg)

        elif msg.get_type() == "RC_CHANNELS":
            self._handle_msg_rc_channels(msg)

    def _on_exit(self):
        logger.debug("mav telemetry connection closed")
        self._mav_connection.close()
        deep_clear(self._mav_connection)
        self._mav_connection = None

        self.ptele_fm_ep = None
        self.ptele_gps_ep = None
        self.ptele_gps_raw_ep = None
        self.ptele_attitude_ep = None
        self.ptele_mission_current_ep = None
        self.ptele_mav_system_time = None
        self.ptele_rc_channels = None

        self._home_position_dict = None
        self._mission_ack_dict = None

    ######################################
    # 直接数据访问接口，直接向自动驾驶仪发送消息
    ######################################
    def send_long_message(self,
                          command: int,
                          confirmation=0,
                          param1=0.0,
                          param2=0.0,
                          param3=0.0,
                          param4=0.0,
                          param5=0.0,
                          param6=0.0,
                          param7=0.0,
                          ignore_dry_run=False):
        """
        透传访问 pymavlink 的长参数消息发送，自动设置目标系统 ID 和目标组件 ID，除 command 其余参数默认为 0
        :param command: 命令代码
        :param confirmation: 是否需要确认
        :param param1: 参数 1
        :param param2: 参数 2
        :param param3: 参数 3
        :param param4: 参数 4
        :param param5: 参数 5
        :param param6: 参数 6
        :param param7: 参数 7
        :param ignore_dry_run: 忽略空跑选项，谨慎启用，可能触发意外的测试问题
        """
        if self._is_dry_run and (not ignore_dry_run):
            return

        self._mav_connection.mav.command_long_send(
            self._mav_connection.target_system,
            self._mav_connection.target_component,
            command,
            confirmation,
            param1, param2, param3, param4, param5, param6, param7
        )

    def send_int_message(self,
                         frame: int,
                         command: int,
                         current=0,
                         autocontinue=0,
                         param1=0.0,
                         param2=0.0,
                         param3=0.0,
                         param4=0.0,
                         x=0,
                         y=0,
                         z=0.0,
                         ignore_dry_run=False):
        """
        透传访问 pymavlink 的整数消息发送，自动设置目标系统 ID 和目标组件 ID，除 command 其余参数默认为 0
        :param frame: 帧类型
        :param command: 命令ID
        :param current: 未使用，设置为0
        :param autocontinue: 未使用，设置为0
        :param param1: 参数1
        :param param2: 参数2
        :param param3: 参数3
        :param param4: 参数4
        :param x: X 整数
        :param y: Y 整数
        :param z: Z 浮点
        :param ignore_dry_run: 忽略空跑选项，谨慎启用，可能触发意外的测试问题
        """
        if self._is_dry_run and (not ignore_dry_run):
            return

        self._mav_connection.mav.command_int_send(
            self._mav_connection.target_system,
            self._mav_connection.target_component,
            frame,
            command,
            current, autocontinue, param1, param2, param3, param4, x, y, z
        )

    def get_home_position(self):
        """
        获取飞行器 HOME 位置，如果 HOME 分析没有稳定或不可用，则返回 None
        """
        if self._home_position_fixed:
            return self._home_position_dict

        return None

    def send_wp_list(self, wp_list: list[list]):
        """
        上传航路点，使用 [[wp_latitude(Deg), wp_longitude(Deg), wp_altitude(Meters)],...]，
        航点发送是一个忙等待函数，不要在实时任务调度中执行。
        :param wp_list: 航点列表，使用起飞目标和降落位置为起点和终点，初始化 HOME 航点自动添加，不要发送
        :return: 航点上传成功返回航点数，失败返回 0
        """
        logger.debug(f"recv way points list {wp_list} try upload to aircraft")
        self._send_base_request(enable=0)

        # 清除任务
        self._mav_connection.mav.mission_clear_all_send(self._mav_connection.target_system,
                                                        self._mav_connection.target_component,
                                                        mavutil.mavlink.MAV_MISSION_TYPE_MISSION)
        logger.debug("aircraft way points list cleared")

        # 使用 MAVProxy 提供的航点加载器添加航点，Ardupilot 要求的航点顺序如下
        # HOME 初始化航点，作为 0 航点，一般与发射点重合
        # TAKEOFF 航点，使用 Ardupilot 预设的自动起飞
        # 中间若干任务航点，最后一个降落航点
        wp = mavwp.MAVWPLoader()
        wp.target_system = self._mav_connection.target_system
        wp.target_component = self._mav_connection.target_component

        wp.add(mavutil.mavlink.MAVLink_mission_item_int_message(
            self._mav_connection.target_system, self._mav_connection.target_component, 0,
            mavutil.mavlink.MAV_FRAME_GLOBAL_RELATIVE_ALT,
            mavutil.mavlink.MAV_CMD_NAV_WAYPOINT,
            0, 0, 0, 0, 0, 0,
            int(self._home_position_dict["latitude"] * 1e7), int(self._home_position_dict["longitude"] * 1e7), 0.0
        ))

        wp.add(mavutil.mavlink.MAVLink_mission_item_int_message(
            self._mav_connection.target_system, self._mav_connection.target_component, 1,
            mavutil.mavlink.MAV_FRAME_GLOBAL_RELATIVE_ALT,
            mavutil.mavlink.MAV_CMD_NAV_TAKEOFF,
            0, 0, 0, 0, 0, 0,
            int(wp_list[0][0] * 1e7), int(wp_list[0][1] * 1e7), float(wp_list[0][2])
        ))

        for wp_i, (lat, lon, alt) in enumerate(wp_list[1:-1]):
            wp.add(mavutil.mavlink.MAVLink_mission_item_int_message(
                self._mav_connection.target_system, self._mav_connection.target_component, wp_i + 2,
                mavutil.mavlink.MAV_FRAME_GLOBAL_RELATIVE_ALT,
                mavutil.mavlink.MAV_CMD_NAV_WAYPOINT,
                0, 0, 0, 0, 0, 0,
                int(lat * 1e7), int(lon * 1e7), float(alt)
            ))

        wp.add(mavutil.mavlink.MAVLink_mission_item_int_message(
            self._mav_connection.target_system, self._mav_connection.target_component, wp.count(),
            mavutil.mavlink.MAV_FRAME_GLOBAL_RELATIVE_ALT,
            mavutil.mavlink.MAV_CMD_NAV_LAND,
            0, 0, 0, 0, 0, 0,
            int(wp_list[-1][0] * 1e7), int(wp_list[-1][1] * 1e7), float(wp_list[-1][2])
        ))

        # 上传任务，这里会忙等待，确保发送流程不被其他通讯打断
        self._mav_connection.waypoint_clear_all_send()
        self._mav_connection.waypoint_count_send(wp.count())
        for i in range(wp.count()):
            msg = self._mav_connection.recv_match(type='MISSION_REQUEST', blocking=True)
            if msg.seq == i:
                self._mav_connection.mav.send(wp.wp(i))
        ack = self._mav_connection.recv_match(type='MISSION_ACK', blocking=True)
        self._send_base_request()

        if ack.type == mavutil.mavlink.MAV_MISSION_ACCEPTED:
            logger.warning("mission uploaded successfully.")
            return wp.count()
        else:
            logger.warning(f"mission upload failed: {ack.type}")
            return 0

    def send_flight_mode(self, fm: int, check=False):
        """
        尝试将飞行模式切换到新模式
        :param fm: 请求的目标飞行模式
        :param check: 启用忙等待执行检查，默认禁用且返回成功执行
        :return: 执行结果，成功 True
        """
        logger.debug(f"require flight mode change to {fm}")

        if check:
            self._send_base_request(enable=0)

        self.send_long_message(mavutil.mavlink.MAV_CMD_DO_SET_MODE, param1=1, param2=fm)

        if check:
            msg = self._mav_connection.recv_match(type='COMMAND_ACK', blocking=True, timeout=1)
            if msg:
                self._send_base_request()

                if msg.command == mavutil.mavlink.MAV_CMD_DO_SET_MODE:
                    if msg.result == mavutil.mavlink.MAV_RESULT_ACCEPTED:
                        return True
                    else:
                        return False
        return True

    def send_mission_current(self, wp_index: int, check=False):
        """
        尝试将当前任务切换新航点
        :param wp_index: 请求的目标飞行模式
        :param check: 启用忙等待执行检查，默认禁用且返回成功执行
        :return: 执行结果，成功 True
        """
        logger.debug(f"require next way point index change to {wp_index}")

        if check:
            self._send_base_request(enable=0)

        # self.mav_connection.mav.command_long_send(
        #     self.mav_connection.target_system,
        #     self.mav_connection.target_component,
        #     mavutil.mavlink.MAV_CMD_DO_SET_MISSION_CURRENT,
        #     0,
        #     wp_index,
        #     0, 0, 0, 0, 0, 0
        # )
        self.send_long_message(mavutil.mavlink.MAV_CMD_DO_SET_MISSION_CURRENT, param1=wp_index)

        if check:
            msg = self._mav_connection.recv_match(type='COMMAND_ACK', blocking=True, timeout=1)
            if msg:
                self._send_base_request()

                if msg.command == mavutil.mavlink.MAV_CMD_DO_SET_MISSION_CURRENT:
                    if msg.result == mavutil.mavlink.MAV_RESULT_ACCEPTED:
                        return True
                    else:
                        return False
        return True

    def send_guided_waypoint(self, lat: float, lon: float, alt: float, velocity=-1.0, l_rad=1, check=False):
        """
        上传 guided 模式航路点，使用 REPOSITION 方法。
        :param lat:
        :param lon:
        :param alt:
        :param velocity:
        :param l_rad:
        :param check:
        :return: 如果启用 check 则成功为 True 失败为 False，禁用则默认 True
        """
        logger.debug(f"sending guided waypoint: lat={lat}, lon={lon}, alt={alt}m")

        if check:
            self._send_base_request(enable=0)

        self.send_int_message(mavutil.mavlink.MAV_FRAME_GLOBAL_RELATIVE_ALT_INT,
                              mavutil.mavlink.MAV_CMD_DO_REPOSITION,
                              param1=velocity,
                              param2=mavutil.mavlink.MAV_DO_REPOSITION_FLAGS_CHANGE_MODE,
                              param3=l_rad, x=int(lat * 1e7), y=int(lon * 1e7), z=float(alt))
        if check:
            msg = self._mav_connection.recv_match(type='COMMAND_ACK', blocking=True, timeout=1)
            if msg:
                self._send_base_request()

                if msg.command == mavutil.mavlink.MAV_CMD_DO_REPOSITION:
                    if msg.result == mavutil.mavlink.MAV_RESULT_ACCEPTED:
                        return True
                    else:
                        return False
        return True

    # TODO 这里使用的数据流请求被 mavlink 组织标记为了 弃用 未来将其升级为更新的 MESSAGE_INTERVAL
    def _send_home_position_request(self):
        self.send_long_message(mavutil.mavlink.MAV_CMD_REQUEST_MESSAGE, param1=242)

    def _send_base_request(self, enable=1):
        # 姿态数据，飞行模式请求
        self._mav_connection.mav.request_data_stream_send(
            self._mav_connection.target_system,
            self._mav_connection.target_component,
            mavutil.mavlink.MAV_DATA_STREAM_ALL,
            self._base_req_rate,
            enable
        )

    def _send_mission_ack_requests(self):
        self.send_long_message(mavutil.mavlink.MAV_CMD_REQUEST_MESSAGE, param1=47)


@handler.add_handler("virtual_bus:plane_tele")
def load_config(veh: Vehicle, cfg_section: dict):
    logger.debug(f"plane telemetry load config {cfg_section}")

    cfg = deepcopy(DEFAULT_CONFIG)
    cfg.update(cfg_section)

    veh.reg_part(cfg["label"], PlaneTeleBus(cfg), is_sche=True)
